import pulp
import numpy as np
from itertools import product
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass


@dataclass
class AlgorithmConfig:
    # 基础参数
    bs_distance: int = 1000  # 基站间距(m)
    system_bandwidth: float = 10e6  # 系统带宽(Hz)
    noise_spectral_density: float = 5e-17  # 噪声功率谱密度(W/Hz)
    max_transmit_power: float = 40.0  # 最大发射功率(W)
    power_amplifier_efficiency: float = 2.8571  # 功率放大器效率
    static_power: float = 10.0  # 静态电路功耗(W)

    # 算法参数
    max_iterations: int = 1000  # 最大迭代次数
    epsilon1: float = 1e-6  # 收敛阈值1
    epsilon2: float = 1e-6  # 收敛阈值2
    max_storage: float = 1000.0  # 最大储能容量(W)
    path_loss_exponent: float = 4.0  # 路径损耗指数
    shadow_fading_std: float = 8.0  # 阴影衰落标准差(dB)
    energy_loss_range: Tuple[float, float] = (0.7, 0.9)  # 能量传输损耗系数范围


class GreenCellNetwork:
    def __init__(self, num_bs: int, num_users: int, area_size: int = 1000,
                 config: AlgorithmConfig = AlgorithmConfig()):
        # 系统参数
        self.num_bs = num_bs
        self.num_users = num_users
        self.area_size = area_size
        self.cfg = config

        # 网络拓扑
        self.bs_coords = self._generate_bs_positions()
        self.user_coords = np.random.uniform(0, area_size, (num_users, 2))

        # 信道模型
        self.g = self._calculate_channel_gain()

        # 能量参数
        self.E = np.random.uniform(0.8 * self.cfg.max_transmit_power,
                                   self.cfg.max_transmit_power, num_bs)
        self.alpha = self._init_energy_loss_factors()

        # 算法状态
        self.X: Optional[Dict[Tuple[int, int], int]] = None
        self.Y: Optional[Dict[Tuple[int, int], float]] = None
        self.P: Optional[np.ndarray] = None
        self.xi: Optional[np.ndarray] = None
        self.convergence_iter: Optional[int] = None
        self.PG_history: List[float] = []

    def _generate_bs_positions(self) -> np.ndarray:
        """生成基站位置（规则六边形布局）"""
        positions = []
        sqrt3 = np.sqrt(3)
        for i in range(self.num_bs):
            x = self.cfg.bs_distance * (i % 3) * 1.5
            y = self.cfg.bs_distance * sqrt3 * (i // 3)
            positions.append([x, y])
        return np.array(positions)

    def _calculate_channel_gain(self) -> np.ndarray:
        """根据论文II-A节计算信道增益"""
        g_matrix = np.zeros((self.num_bs, self.num_users))

        for n in range(self.num_bs):
            for m in range(self.num_users):
                # 计算距离
                dx = self.bs_coords[n, 0] - self.user_coords[m, 0]
                dy = self.bs_coords[n, 1] - self.user_coords[m, 1]
                distance = max(np.sqrt(dx ** 2 + dy ** 2), 1e-3)  # 避免除零

                # 路径损耗模型（3GPP模型）
                PL = 128.1 + 37.6 * np.log10(distance / 1000)  # dB

                # 阴影衰落
                shadowing = np.random.normal(0, self.cfg.shadow_fading_std)

                # 综合信道增益（线性值）
                g_matrix[n, m] = 10 ** (-(PL + shadowing) / 10)

        return g_matrix

    def _init_energy_loss_factors(self) -> np.ndarray:
        """初始化能量传输损耗系数矩阵"""
        alpha = np.random.uniform(*self.cfg.energy_loss_range,
                                  (self.num_bs, self.num_bs))
        np.fill_diagonal(alpha, 0)  # 自身传输损耗为0
        return alpha

    def _solve_LDP(self) -> Tuple[Dict, Dict]:
        # 初始化参数
        BL = {n: set() for n in range(self.num_bs)}
        X_prime = {n: {m: 0 for m in range(self.num_users)} for n in range(self.num_bs)}
        R_bar = self._calculate_R_bar()
        R_req = np.full(self.num_users, 3e6)  # 假设统一速率需求

        iteration = 0
        while iteration < self.cfg.max_iterations:
            # 步骤3-6：求解松弛问题
            prob = pulp.LpProblem("LDP_Relaxation", pulp.LpMaximize)
            x = pulp.LpVariable.dicts("x", [(n, m) for n, m in product(range(self.num_bs), range(self.num_users))],
                                      lowBound=0, upBound=1)
            phi = pulp.LpVariable("phi", lowBound=0)

            # 目标函数
            prob += phi

            # 约束条件
            for m in range(self.num_users):
                prob += pulp.lpSum(x[(n, m)] * R_bar[n, m] for n in range(self.num_bs)) >= R_req[m] + phi

            for n in range(self.num_bs):
                prob += pulp.lpSum(x[(n, m)] for m in range(self.num_users)) <= 1
                for m in BL[n]:
                    prob += x[(n, m)] == 0

            # 求解
            status = prob.solve(pulp.PULP_CBC_CMD(msg=False))

            if status != 1:  # 无可行解
                return None, None

            # 步骤7-13：黑名单更新
            X_current = {(n, m): pulp.value(x[(n, m)]) for n, m in product(range(self.num_bs), range(self.num_users))}
            omega = 1

            for n in range(self.num_bs):
                total_load = sum(X_current[(n, m)] * R_req[m] / R_bar[n, m] for m in range(self.num_users))
                if total_load > 1 + 1e-6:
                    omega = 0
                    newly_added = [m for m in range(self.num_users)
                                   if X_current[(n, m)] > 0.9 and X_prime[n][m] < 0.1]
                    if newly_added:
                        m_star = max(newly_added, key=lambda m: X_current[(n, m)] * R_req[m] / R_bar[n, m])
                        BL[n].add(m_star)

            if omega == 1:
                break

            X_prime = X_current.copy()
            iteration += 1

        # 步骤20：资源分配
        Y = {}
        for n in range(self.num_bs):
            total_resource = sum(X_current[(n, m)] * R_req[m] / R_bar[n, m] for m in range(self.num_users))
            for m in range(self.num_users):
                if X_current[(n, m)] > 0.9:
                    Y[(n, m)] = (R_req[m] / R_bar[n, m] +
                                 (1 - total_resource) / sum(1 / R_bar[n, m_prime]
                                                            for m_prime in range(self.num_users) if
                                                            X_current[(n, m_prime)] > 0.9))
                else:
                    Y[(n, m)] = 0

        return X_current, Y

    def _energy_cooperation(self, P: np.ndarray) -> Tuple[np.ndarray, int]:
        # 计算能量需求
        energy_demand = self.cfg.power_amplifier_efficiency * P + self.cfg.static_power
        surplus_mask = energy_demand < self.E
        deficit_mask = ~surplus_mask

        # 能量收集
        surplus_energy = self.E[surplus_mask] - energy_demand[surplus_mask]
        collected_energy = np.sum(surplus_energy * self.alpha[surplus_mask, 0])  # 假设aggregator索引为0

        # 能量分配
        xi = self.E.copy()
        sorted_deficit = np.argsort(-self.alpha[0, deficit_mask])

        for idx in sorted_deficit:
            required = energy_demand[idx] - self.E[idx]
            available = min(collected_energy / self.alpha[0, idx], required)
            xi[idx] += available * self.alpha[0, idx]
            collected_energy -= available

        # 双向传输约束
        for i in range(self.num_bs):
            for j in range(i + 1, self.num_bs):
                if xi[i] > self.E[i] and xi[j] > self.E[j]:
                    if self.alpha[i, j] > self.alpha[j, i]:
                        xi[j] = self.E[j]
                    else:
                        xi[i] = self.E[i]

        zeta = 1 if np.all(energy_demand <= xi) else 0
        return xi, zeta

    def run(self) -> Dict:
        """主算法入口"""
        try:
            # 初始化
            P = np.full(self.num_bs, self.cfg.max_transmit_power)
            PG = np.sum(np.maximum(
                self.cfg.power_amplifier_efficiency * P + self.cfg.static_power - self.E, 0))
            self.PG_history = [PG]

            for t in range(1, self.cfg.max_iterations + 1):
                # 负载分布
                X, Y = self._solve_LDP()
                if X is None:
                    raise ValueError("No feasible solution in LDP")

                # 功率控制
                P_new = self._power_control(X, Y)

                # 能量合作
                xi, zeta = self._energy_cooperation(P_new)

                # 收敛判断
                PG_new = np.sum(np.maximum(
                    self.cfg.power_amplifier_efficiency * P_new + self.cfg.static_power - xi, 0))

                delta_P = np.linalg.norm(P_new - P)
                delta_PG = abs(PG_new - PG)

                # 更新状态
                P = P_new
                PG = PG_new
                self.PG_history.append(PG)

                # 双重收敛条件
                if (zeta == 1 or
                        (delta_PG < self.cfg.epsilon1 and
                         delta_P < self.cfg.epsilon2)):
                    self.convergence_iter = t
                    break

            # 保存最终状态
            self.X = X
            self.Y = Y
            self.P = P
            self.xi = xi

            return self._format_results()

        except Exception as e:
            return {
                'status': 'error',
                'message': str(e),
                'convergence': False
            }

    def _format_results(self) -> Dict:
        """格式化输出结果"""
        # 基站信息
        stations = []
        for n in range(self.num_bs):
            stations.append({
                'id': n,
                'position': self.bs_coords[n].tolist(),
                'transmit_power': float(self.P[n]),
                'energy_supply': float(self.xi[n]),
                'connected_users': [
                    m for m in range(self.num_users)
                    if self.X.get((n, m), 0) > 0.9
                ]
            })

        # 用户信息
        users = []
        for m in range(self.num_users):
            connected_bs = next(
                (n for n in range(self.num_bs) if self.X.get((n, m), 0) > 0.9),
                None)
            users.append({
                'id': m,
                'position': self.user_coords[m].tolist(),
                'connected_bs': connected_bs,
                'achieved_rate': float(
                    self.Y.get((connected_bs, m), 0) *
                    self.cfg.system_bandwidth *
                    np.log2(1 + self.P[connected_bs] * self.g[connected_bs, m] /
                            (self.cfg.noise_spectral_density * self.cfg.system_bandwidth))
                ) if connected_bs is not None else 0
            })

        # 全局指标
        metrics = {
            'total_grid_power': float(self.PG_history[-1]),
            'iteration_steps': self.convergence_iter or len(self.PG_history),
            'power_history': [float(p) for p in self.PG_history],
            'energy_efficiency': float(
                sum(u['achieved_rate'] for u in users) /
                (self.PG_history[-1] + 1e-6))  # 防止除零
        }

        return {
            'status': 'success',
            'stations': stations,
            'users': users,
            'metrics': metrics
        }